﻿using System;
using System.Collections.Generic;
using System.Collections.ObjectModel;
using System.Linq;
using System.ServiceModel;
using System.ServiceModel.Channels;
using System.ServiceModel.Dispatcher;
using System.Text;

namespace NhJump.Persistence.Context.WcfBinding
{
    public class WcfUnitOfWorkMessageInspector : IDispatchMessageInspector
    {
        private readonly Type _serviceType;

        public WcfUnitOfWorkMessageInspector(Type serviceType)
        {
            _serviceType = serviceType;
        }

        public object AfterReceiveRequest(ref Message request, IClientChannel channel, InstanceContext instanceContext)
        {
            instanceContext.Extensions.Add(new WcfUnitOfWorkExtension(UnitOfWorkFactory.ForServiceType(_serviceType)));
            return null;
        }

        public void BeforeSendReply(ref Message reply, object correlationState)
        {
            // we no longer need the extension or error handlers

            var extensions = OperationContext.Current.InstanceContext.Extensions.FindAll<WcfUnitOfWorkExtension>();

            foreach (var extension in extensions)
            {
                OperationContext.Current.InstanceContext.Extensions.Remove(extension);
            }

            var errorHandlers =
                new List<IErrorHandler>(OperationContext.Current.EndpointDispatcher.ChannelDispatcher.ErrorHandlers.
                                            Where(
                                            h => h.GetType() == typeof(WcfUnitOfWorkErrorHandler)));

            foreach (var errorHandler in errorHandlers)
            {
                OperationContext.Current.EndpointDispatcher.ChannelDispatcher.ErrorHandlers.Remove(errorHandler);
            }
        }
    }
}